import sys
import pandas as pd
import numpy as np
from PyQt5.QtWidgets import QApplication
import pyqtgraph as pg

# ----------------------
# 示例数据
# ----------------------
# 假设你的 DataFrame 格式如下：
# trade_date, open, high, low, close, pre_close, change, pct_chg, vol, amount

# df = pd.DataFrame({
#     'trade_date': pd.date_range('2025-01-01', periods=100),
#     'open': np.random.rand(100)*10+100,
#     'high': np.random.rand(100)*10+105,
#     'low': np.random.rand(100)*10+95,
#     'close': np.random.rand(100)*10+100,
#     'vol': np.random.randint(1000, 5000, 100)
# })

# 读取数据
df = pd.read_csv("F:\\股票数据\\daily20250829\\daily\\000554.SZ.csv")
# 转换日期格式
df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
# 按日期排序
df = df.sort_values('trade_date').reset_index(drop=True)

print(df.head())


# ----------------------
# K 线转换函数
# ----------------------
def make_candlestick_data(df):
    """把 DataFrame 转成 PyQtGraph 能绘制的 K 线数据"""
    candlesticks = []
    for i, row in df.iterrows():
        t = i  # x轴用索引
        open_price = row['open']
        high_price = row['high']
        low_price = row['low']
        close_price = row['close']
        candlesticks.append((t, open_price, close_price, low_price, high_price))
    return candlesticks

candlesticks = make_candlestick_data(df)

# ----------------------
# PyQtGraph 窗口
# ----------------------
app = QApplication(sys.argv)
win = pg.GraphicsLayoutWidget(show=True, title="K线图与成交量 - PyQtGraph")
win.resize(1000,800)

# 创建上下两个plot
# 上面的K线图
kline_plot = win.addPlot(title="K线图", row=0, col=0)
kline_plot.showGrid(x=True, y=True)

# 下面的成交量图
volume_plot = win.addPlot(title="成交量", row=1, col=0)
volume_plot.showGrid(x=True, y=True)

# 禁用自动范围调整，防止鼠标移动时K线图高度变化
kline_plot.vb.disableAutoRange()
volume_plot.vb.disableAutoRange()

# 创建文本标签用于显示鼠标悬停信息
kline_info_text = pg.TextItem(anchor=(0, 0), color='w', fill=pg.mkColor(0, 0, 0, 200))
kline_plot.addItem(kline_info_text)
kline_info_text.hide()
kline_info_text.setZValue(100)  # 设置最高层级，确保显示在最前面

volume_info_text = pg.TextItem(anchor=(0, 0), color='w', fill=pg.mkColor(0, 0, 0, 200))
volume_plot.addItem(volume_info_text)
volume_info_text.hide()
volume_info_text.setZValue(100)  # 设置最高层级，确保显示在最前面

# ----------------------
# 自定义 K 线 Item
# ----------------------
class CandlestickItem(pg.GraphicsObject):
    def __init__(self, data):
        pg.GraphicsObject.__init__(self)
        self.data = data
        self.generatePicture()

    def generatePicture(self):
        self.picture = pg.QtGui.QPicture()
        p = pg.QtGui.QPainter(self.picture)
        w = 0.4  # K线宽度
        for (t, open, close, low, high) in self.data:
            if close >= open:
                p.setPen(pg.mkPen('r'))  # 涨
                p.setBrush(pg.mkBrush('r'))
            else:
                p.setPen(pg.mkPen('g'))  # 跌
                p.setBrush(pg.mkBrush('g'))
            # 画实体
            p.drawRect(pg.QtCore.QRectF(t-w, open, w*2, close-open))
            # 画上下影线
            p.drawLine(pg.QtCore.QPointF(t, low), pg.QtCore.QPointF(t, high))
        p.end()

    def paint(self, painter, option, widget):
        painter.drawPicture(0,0,self.picture)

    def boundingRect(self):
        return pg.QtCore.QRectF(self.picture.boundingRect())

# 添加 K 线到上面的plot
item = CandlestickItem(candlesticks)
kline_plot.addItem(item)

# 设置K线图的固定Y轴范围
min_price = min(df['low'])
max_price = max(df['high'])
price_range = max_price - min_price
kline_plot.setYRange(min_price - price_range * 0.1, max_price + price_range * 0.1)

# ----------------------
# 成交量图
# ----------------------
# 创建成交量数据，根据涨跌设置颜色
volume_x = []
volume_heights = []
volume_colors = []

for i, row in df.iterrows():
    volume_x.append(i)
    volume_heights.append(row['vol'])
    # 根据收盘价与开盘价比较设置颜色
    if row['close'] >= row['open']:
        volume_colors.append('r')  # 涨 - 绿色
    else:
        volume_colors.append('g')  # 跌 - 红色

# 绘制成交量柱状图
# 将成交量高度按比例缩小，避免显示过高
max_volume = max(volume_heights)
min_volume = min(volume_heights)
# scaled_heights = [h * 0.6 for h in volume_heights]  # 缩放到60%
# 归一化并按K线图比例缩放
scaled_heights = [h / max(volume_heights) for h in volume_heights]  # 归一化到 [0,1]
scaled_heights = [h * 1000 * kline_plot.viewRange()[1][1] for h in scaled_heights]  # 按比例缩放

volume_bars = pg.BarGraphItem(x=volume_x, 
                             height=scaled_heights, 
                             width=0.8,
                             brushes=volume_colors,
                             pens=volume_colors)
volume_plot.addItem(volume_bars)

# 设置成交量图的Y轴范围，压缩高低差
# 使用较小的Y轴范围来减少视觉上的高低差异
volume_range = max_volume - min_volume
volume_plot.setYRange(min_volume, min_volume + volume_range * 0.2, padding=0.05)

# 设置两个图的X轴范围
kline_plot.setXRange(0, len(df) - 1)
volume_plot.setXRange(0, len(df) - 1)

# 移除X轴同步缩放，让两个图独立缩放
kline_plot.setXLink(volume_plot)  # 注释掉这行

# ----------------------
# X轴显示日期
# ----------------------
xticks = [(i, str(d.date())) for i, d in enumerate(df['trade_date'])]

# 为K线图设置X轴
kline_axis = kline_plot.getAxis('bottom')
kline_axis.setTicks([xticks[::max(1,len(xticks)//10)]])

# 为成交量图设置X轴
volume_axis = volume_plot.getAxis('bottom')
volume_axis.setTicks([xticks[::max(1,len(xticks)//10)]])

# ----------------------
# 鼠标悬停显示信息功能
# ----------------------
def on_kline_mouse_move(signal):
    """K线图鼠标移动事件"""
    try:
        pos = kline_plot.vb.mapSceneToView(signal[0])
        x_pos = int(round(pos.x()))
        
        # 检查是否在有效范围内
        if 0 <= x_pos < len(df):
            row = df.iloc[x_pos]
            info = f"""日期: {row['trade_date'].strftime('%Y-%m-%d')}
开盘: {row['open']:.2f}
最高: {row['high']:.2f}
最低: {row['low']:.2f}
收盘: {row['close']:.2f}
涨跌: {row['close'] - row['open']:.2f}
涨跌幅: {((row['close'] - row['open']) / row['open'] * 100):.2f}%"""
            
            kline_info_text.setText(info)
            kline_info_text.setPos(pos.x(), pos.y())
            kline_info_text.show()
        else:
            kline_info_text.hide()
    except Exception as e:
        print(f"K线图鼠标事件错误: {e}")
        kline_info_text.hide()

def on_volume_mouse_move(signal):
    """成交量图鼠标移动事件"""
    try:
        pos = volume_plot.vb.mapSceneToView(signal[0])
        x_pos = int(round(pos.x()))
        
        # 检查是否在有效范围内
        if 0 <= x_pos < len(df):
            row = df.iloc[x_pos]
            info = f"""日期: {row['trade_date'].strftime('%Y-%m-%d')}
成交量: {row['vol']:,} 手
成交额: {row['amount']:,.0f} 元"""
            
            volume_info_text.setText(info)
            volume_info_text.setPos(pos.x(), pos.y())
            volume_info_text.show()
        else:
            volume_info_text.hide()
    except Exception as e:
        print(f"成交量图鼠标事件错误: {e}")
        volume_info_text.hide()

# 连接鼠标事件
kline_proxy = pg.SignalProxy(kline_plot.scene().sigMouseMoved, rateLimit=60, slot=on_kline_mouse_move)
volume_proxy = pg.SignalProxy(volume_plot.scene().sigMouseMoved, rateLimit=60, slot=on_volume_mouse_move)

sys.exit(app.exec_())
